import asyncio
from typing import TypedDict

from langgraph.constants import START
from langgraph.graph import StateGraph, MessagesState, END


class Mystate(TypedDict):
    aaa: str


# 1. 定义节点函数
async def step1(state: Mystate):
    for i in range(5):
        print(i)
    state["aaa"] = "我来了step1"
    return state


async def step2(state: Mystate):
    state["aaa"] = "我来了step2"
    return state


# 2. 构建Graph
graph = StateGraph(MessagesState)
graph.add_node("step1", step1)
graph.add_node("step2", step2)
graph.add_edge(START, "step1")
graph.add_edge("step1", "step2")
graph.add_edge("step2", END)
app = graph.compile()


# 3. 异步astream流式输出
async def run_astream():
    async for event in app.astream({}):  #
        print("---- AStream Event ----")
        print(event)


if __name__ == "__main__":
    asyncio.run(run_astream())
